package grpc_proxy_middleware

import (
	"fmt"
	"gitee.com/zhoulvvv/my_go_gateway/dao"
	"gitee.com/zhoulvvv/my_go_gateway/public"
	"github.com/pkg/errors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/peer"
	"log"
	"strings"
)

// GRPC IP 黑名单中间件
func GRPCBlackListModeMiddleware(service *dao.ServiceDetail) func(srv interface{}, ss grpc.ServerStream,
	info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		// 获取IP白名单集合
		ipWhileList := []string{}
		if service.AccessControl.WhiteList != "" {
			ipWhileList = strings.Split(service.AccessControl.WhiteList, ",")
		}
		// 获取IP黑名单集合
		ipBlackList := []string{}
		if service.AccessControl.BlackList != "" {
			ipBlackList = strings.Split(service.AccessControl.BlackList, ",")
		}
		peerCtx, ok := peer.FromContext(ss.Context())
		if !ok {
			log.Printf("【GRPC IP 黑名单中间件中获取客户端IP失败】peer.FromContext(ss.Context()) failed with error")
		}
		splits := strings.Split(peerCtx.Addr.String(), ":")
		clientIP := ""
		if len(splits) == 2 {
			clientIP = splits[0]
		}
		// 判断是否开启权限校验
		if service.AccessControl.OpenAuth == 1 && len(ipWhileList) == 0 && len(ipBlackList) > 0 {
			if public.InStringsSlice(ipBlackList, clientIP) {
				return errors.New(fmt.Sprintf("%s IP 没有权限访问", clientIP))
			}
		}
		if err := handler(srv, ss); err != nil {
			log.Printf("【GRPC】RPC failed with error %v\n", err)
		}
		return nil
	}
}
